package cn.lsh.kafka.listener;

import cn.lsh.kafka.db.entity.OffsetEntity;
import cn.lsh.kafka.db.service.OffsetService;
import cn.lsh.kafka.db.service.ServiceFactory;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.common.TopicPartition;

import java.util.Collection;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;

/**
 * 自定义存储offset时，实现自己的Rebalance监听器
 */
public class MyConsumerRebalanceListener implements ConsumerRebalanceListener {
	private Consumer consumer;

	public static final Map<TopicPartition, Long> currentOffset = new HashMap<>();

	public MyConsumerRebalanceListener(Consumer consumer) {
		this.consumer = consumer;
	}

	@Override
	public void onPartitionsRevoked(Collection<TopicPartition> topicPartitions) {
		//方法在Rebalance之前调用，将当前消费者内存中保存的offset信息保存到自定义的存储中
		// OffsetService service = new OffsetServiceImpl();
		OffsetService service = ServiceFactory.buildService(OffsetService.class);
		commitOffset(consumer.groupMetadata().groupId(), service);
		// service.commit();
		// service.colse();
	}

	@Override
	public void onPartitionsAssigned(Collection<TopicPartition> topicPartitions) {
		System.out.println("Rebalance拉取offset");
		String groupId = consumer.groupMetadata().groupId();
		//方法在Rebalance之后调用
		currentOffset.clear();//清空客户端内存中的offset
		for (TopicPartition partition : topicPartitions) {
			//定位到最近提交的offset位置继续消费
			consumer.seek(partition, getOffset(groupId, partition));
		}
	}

	/**
	 * 自定义实现获取offset
	 *
	 * @param partition
	 * @return
	 */
	private static long getOffset(String groupId, TopicPartition partition) {
		OffsetService service = ServiceFactory.buildService(OffsetService.class);
		OffsetEntity entity = service.queryByCond(groupId, partition.topic(), partition.partition());
		if (entity != null) {
			return entity.getValue();
		}
		return 0;
	}

	/**
	 * 自定义实现保存offset
	 */
	public static void commitOffset(String groupId, OffsetService service) {
		Date now = new Date();
		for (Map.Entry<TopicPartition, Long> entry : currentOffset.entrySet()) {
			TopicPartition partition = entry.getKey();
			OffsetEntity entity = service.queryByCond(groupId, partition.topic(), partition.partition());
			//注意：消费者保存到currentOffset的offset是当前消费消息的偏移量，这里提交到mysql需要加1,不然Rebalance后会重复消费一条消息
			long nextOffset = entry.getValue() + 1;
			if (entity == null) {
				OffsetEntity offset = new OffsetEntity();
				offset.setGroupId(groupId);
				offset.setTopic(partition.topic());
				offset.setPartitionId(partition.partition());
				offset.setValue(nextOffset);
				offset.setCreateDate(now);
				offset.setUpdateDate(now);
				service.insert(offset);
			} else if (nextOffset != entity.getValue()) {
				service.update(entity.getId(), nextOffset);
			}
		}
	}
}
